# 年末检视持仓
# 1. 查看pr > 0.7 的股票
# 1.1 如果有的话就卖出
# 1.1.1 卖出后寻找符合条件的新票并买入

import baostock as bs
import pandas as pd
import matplotlib.pyplot as plt
import pickle

from utils.TimeUtils import get_last_workday_of_year
from utils.methods import get_pe, get_roe, get_close_price, to_rate, get_stable_code_from_hs300


class PerformanceTest:
    def __init__(self, initial_codes, start_date, end_date):
        self.holdings = {}
        self.start_date = start_date
        self.end_date = end_date
        self.start_year = int(start_date[:4])
        self.end_year = int(end_date[:4])
        self.trade_log = {}

        self.cache_stable_code = {}

        # 获取start_date的股价
        print('获取start_date的股价')
        code_dict = {}
        open_sum = 0.0
        for code in initial_codes:
            rs = bs.query_history_k_data_plus(code, "code,open", start_date=self.start_date,
                                              end_date=self.end_date, frequency="d", adjustflag="3")
            data_list = []
            while (rs.error_code == '0') & rs.next():
                # 获取一条记录，将记录合并在一起
                data_list.append(rs.get_row_data())
            result = pd.DataFrame(data_list, columns=rs.fields)
            init_price = float(result['open'][0])
            open_sum += init_price
            code_dict[code] = {
                'init_price': init_price,
            }
        # 建仓
        print('建仓')
        i = 1
        for code in initial_codes:
            self.holdings[code] = {
                'amount': int(open_sum / code_dict[code]['init_price'] * 100),
                'code': code,
                'cost_price': code_dict[code]['init_price']
            }
            print(
                f'{i}/{len(initial_codes)} {code} ({self.holdings[code]["cost_price"]}/{self.holdings[code]["amount"]})')
            i += 1

    def test(self):
        print('测试')
        for year in range(self.start_year, self.end_year + 1):
            print(f'======== {year} ========')
            last_workday_of_year = str(get_last_workday_of_year(year))
            # 打印输出持仓收益率
            rates = []
            for code in self.holdings:
                if self.holdings[code]['amount'] > 0:
                    # 获取当前价格
                    try:
                        close_price = get_close_price(code, last_workday_of_year)
                        cost_price = self.holdings[code]['cost_price']
                        profit_rate = (close_price - cost_price) / cost_price * 100
                        rates.append([code, profit_rate])
                    except Exception as e:
                        print(e)
            # 按收益率降序排序
            rates.sort(key=lambda x: x[1], reverse=True)
            print('持仓收益率')
            for code, profit_rate in rates:
                print(f'{code} | {profit_rate:.2f}%')
            print(f'{year}年持仓的平均收益率: {(sum([x[1] for x in rates]) / len(rates)):.2f}%')
            print('====================')
            # 检查pr
            for code in list(self.holdings):
                try:
                    pe = get_pe(code, last_workday_of_year)
                    avg_roe = get_roe(code, year - 4, year)
                    pr = pe / (avg_roe * 100)
                except Exception as e:
                    print(e)
                    print(f'{code} {last_workday_of_year} 计算失败')
                    continue
                if pr >= 0.7 and self.holdings[code]['amount'] > 0:
                    # 卖出
                    close_price = get_close_price(code, last_workday_of_year)
                    sold_amount = self.holdings[code]['amount']
                    sold_cash = sold_amount * close_price
                    print(f'{year} {code} 卖出')
                    self.holdings[code]['amount'] = 0
                    trade = {
                        'date': last_workday_of_year,
                        'code': code,
                        'type': 'sell',
                        'amount': sold_amount,
                        'price': get_close_price(code, last_workday_of_year),
                    }
                    if last_workday_of_year in self.trade_log:
                        self.trade_log[last_workday_of_year].append(trade)
                    else:
                        self.trade_log[last_workday_of_year] = [trade]
                    # 买入新的
                    # 获取符合条件的股票
                    if year not in self.cache_stable_code:
                        self.cache_stable_code[year] = get_stable_code_from_hs300(year)
                    # 计算市赚率，按市赚率升序排序
                    result = pd.DataFrame(columns=['code', 'code_name', 'pr', 'pe', 'avg_roe'])
                    for index, stable_code in self.cache_stable_code[year].iterrows():
                        code = stable_code['code']
                        pe = get_pe(code, last_workday_of_year)
                        avg_roe = get_roe(code, year - 4, year)
                        result.loc[len(result)] = [code, stable_code['code_name'], pe / (avg_roe * 100), pe,
                                                   avg_roe]
                    result = result.sort_values(by='pr', ascending=True)
                    # 过滤已有的持仓
                    result = result[~result['code'].isin(self.holdings)]
                    # 购买
                    if len(result) > 0:
                        code = result.iloc[0]['code']
                        # 获取当天的收盘价
                        close_price = get_close_price(code, last_workday_of_year)
                        amount = int(sold_cash / close_price)
                        print(f'{year} {code} 买入')
                        self.holdings[code] = {
                            'amount': amount,
                            'code': code,
                            'cost_price': close_price
                        }
                        trade = {
                            'date': last_workday_of_year,
                            'code': code,
                            'type': 'buy',
                            'amount': amount,
                            'price': close_price,
                        }
                        self.trade_log[last_workday_of_year].append(trade)

    def plot_profit_rate(self):
        # 拉取k线数据
        k_data_dict = {}
        print('获取k线数据')
        index = 1
        for code in self.holdings:
            print(f'{index}. {code}')
            index += 1
            rs = bs.query_history_k_data_plus(code,
                                              "date,open,close",
                                              start_date=self.start_date, end_date=self.end_date,
                                              frequency="d", adjustflag="3")
            data_list = []
            while (rs.error_code == '0') & rs.next():
                # 获取一条记录，将记录合并在一起
                data_list.append(rs.get_row_data())
            result = pd.DataFrame(data_list, columns=rs.fields)
            k_data_dict[code] = {}
            k_data_dict[code]['k_data'] = result
        amount_log = {}
        for code in self.holdings:
            k_data = k_data_dict[code]
            for index, row in k_data['k_data'].iterrows():
                date = row['date']
                if date not in amount_log:
                    amount_log[date] = 0
                # 处理交易
                if date in self.trade_log:
                    for trade in self.trade_log[date]:
                        if trade['code'] == code:
                            if trade['type'] == 'buy':
                                amount_log[date] += trade['amount'] * trade['price']
                            elif trade['type'] == 'sell':
                                amount_log[date] -= trade['amount'] * trade['price']
                if date in amount_log:
                    amount_log[date] += float(row['close']) * self.holdings[code]['amount']
                else:
                    amount_log[date] = float(row['close']) * self.holdings[code]['amount']
        # 转换为list并按date升序排序
        date_value_list = []
        for date, value in amount_log.items():
            date_value_list.append({'date': date, 'value': value})
        date_value_list.sort(key=lambda x: x['date'])
        # 转换为收益率
        profit_rate = to_rate(date_value_list)
        # 转换为DataFrame
        df = pd.DataFrame(profit_rate)
        # 打印输出最终的收益率
        print(f'策略收益率为：{df.iloc[-1]["rate"] * 100.0:.2f}%')
        # 获取沪深300数据
        hs300_data = bs.query_history_k_data_plus("sh.000300",
                                                  "date,code,open,high,low,close,preclose,volume,amount,pctChg",
                                                  start_date=self.start_date, end_date=self.end_date, frequency="d")
        data_list = []
        while (hs300_data.error_code == '0') & hs300_data.next():
            # 获取一条记录，将记录合并在一起
            data_list.append(hs300_data.get_row_data())
        result = pd.DataFrame(data_list, columns=hs300_data.fields)
        # 转换为收益率
        hs300_date_list = result[['date', 'close']].values.tolist()
        hs300_date_list = list(map(lambda r: {'date': r[0], 'value': float(r[1])}, hs300_date_list))
        hs300_rate = pd.DataFrame(to_rate(hs300_date_list))
        print(f'沪深300收益率为：{hs300_rate.iloc[-1]["rate"] * 100.0:.2f}%')
        # 绘制折线图
        plt.plot(df['date'], df['rate'], label='Strategy 1')
        plt.plot(hs300_rate['date'], hs300_rate['rate'], label='HS300')
        plt.xlabel('Date')
        plt.ylabel('Profit Rate(%)')
        plt.legend()
        plt.title(f'Historical Performance {self.start_date} - {self.end_date}')
        plt.show()


if __name__ == '__main__':
    # 回测
    # codes = pd.read_csv('~/Desktop/low_pr.csv')['code'].to_list()
    bs.login()
    # pt = PerformanceTest(codes, '2016-01-01', '2024-12-31')
    # pt.test()
    # with open('low_pr.pkl', 'wb') as f:
    #     pickle.dump(pt, f, pickle.HIGHEST_PROTOCOL)
    # 绘制折线图
    with open('strategy1.pkl', 'rb') as f:
        pt = pickle.load(f)
    pt.plot_profit_rate()
    bs.logout()
